import os
import gymnasium as gym
from train import REINFORCE_Agent
import matplotlib.pyplot as plt
from config import (GAME_NAME, HIDDEN_DIM, DEVICE, LEARNING_RATE,DISCOUNT_FACTOR)

env = gym.make(GAME_NAME, render_mode=None, max_episode_steps=500)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
action_low = env.action_space.low
action_high = env.action_space.high

# 创建智能体
agent = REINFORCE_Agent(
    state_dim=state_dim,
    hidden_dim=HIDDEN_DIM,
    action_dim=action_dim,
    device=DEVICE,
    learning_rate=LEARNING_RATE,
    discount_factor=DISCOUNT_FACTOR,
    action_low=action_low,
    action_high=action_high,
)
# 加载已有模型
results_dir = os.path.join("results", GAME_NAME)
os.makedirs(results_dir, exist_ok=True)
model_path = os.path.join(results_dir, "reinforce_model.pth")
agent.load_model(model_path)

total_returns = []
for ep in range(100):
    episode_return = 0.0
    done = False
    state, _ = env.reset()
    while not done:
        action, _ = agent.take_action(state)
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        state = next_state
        episode_return += reward
    total_returns.append(episode_return)

env.close()

plt.figure(figsize=(10, 6))
plt.plot(range(1, len(total_returns) + 1), total_returns, label='Episode Return', color='steelblue')
plt.xlabel('Episode')
plt.ylabel('Return')
plt.title(f'Test Performance of REINFORCE on {GAME_NAME}')
plt.grid(True)
plt.legend()
plot_path = os.path.join(results_dir, "test_returns.png")
plt.savefig(plot_path)
plt.show()